# import struct

import torch
# import math
import concurrent.futures
import time
# from multiprocessing import Pool, Manager
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
from tqdm import tqdm, trange
import torch.nn.functional as F
# import torch.multiprocessing as mp

# import multiprocessing_on_dill.dummy as mp
# from torchvision.utils import save_image
# import multiprocessing as mp
import cv2
import numpy as np
from PIL import Image
# from threading import Thread
# from cProfile import Profile
# import psutil
import ctypes
from ctypes import c_int, c_void_p, c_float, c_bool
import math
import gc



BLOCK_X = 16
BLOCK_Y = 16
BLOCK_SIZE = BLOCK_X * BLOCK_Y

SH_C0 = 0.28209479177387814
SH_C1 = 0.4886025119029199
SH_C2 = torch.Tensor([1.0925484305920792,
                      -1.0925484305920792,
                      0.31539156525252005,
                      -1.0925484305920792,
                      0.5462742152960396])

SH_C3 = torch.Tensor([-0.5900435899266435,
                      2.890611442640554,
                      -0.4570457994644658,
                      0.3731763325901154,
                      -0.4570457994644658,
                      1.445305721320277,
                      -0.5900435899266435])

lib = ctypes.CDLL("build/libtest.dll")
GetRanges_tile = lib.GetRanges_tile
GetRanges_tile.restype = None
GetRanges_tile.argtypes = [c_int, c_void_p, c_void_p, c_void_p, c_void_p, c_int, c_void_p, c_void_p, c_void_p, c_void_p]

GetRanges_tile2 = lib.GetRanges_tile2
GetRanges_tile2.restype = None
GetRanges_tile2.argtypes = [c_int, c_void_p, c_void_p, c_void_p, c_void_p, c_int, c_void_p, c_void_p, c_void_p, c_void_p]

render_tile = lib.render_tile
render_tile.restype = None
render_tile.argtypes = [c_int, c_void_p, c_void_p, c_void_p, c_void_p, c_void_p, c_void_p, c_int, c_void_p, c_int, c_int]

preprocess = lib.preprocess
preprocess.restype = None
preprocess.argtypes = [c_int,c_int,c_int,c_int,c_void_p,c_void_p,c_float,c_void_p,c_void_p,c_void_p,c_void_p,c_void_p,c_void_p,c_void_p,c_void_p,c_void_p,c_int,c_int,c_float,c_float,c_float,c_float,c_void_p,c_void_p,c_void_p,c_void_p,c_void_p,c_void_p,c_void_p,c_void_p,c_bool,c_void_p,c_void_p]

duplicateWithKeys = lib.duplicateWithKeys
duplicateWithKeys.restype = None
duplicateWithKeys.argtypes = [c_int, c_void_p, c_void_p,c_void_p,c_void_p,c_void_p,c_void_p,c_void_p,]

GetTileRanges = lib.GetTileRanges
GetTileRanges.restype = None
GetTileRanges.argtypes = [c_int, c_void_p, c_void_p]

render = lib.render
render.restype = None
render.argtypes = [c_void_p,c_void_p,c_void_p,c_void_p,c_void_p,c_void_p,c_void_p,c_int,c_int]


class MyRaster:
    def __init__(self, w, h):
        self.image = torch.zeros((w, h, 3))
        self.total = 0
        self.ranges = []
        self.non_empty_tile_id = []
        self.rect_min = None
        self.rect_max = None
        x = torch.Tensor(range(0, BLOCK_X)).int().view(BLOCK_X, 1)
        y = torch.Tensor(range(0, BLOCK_Y)).int().view(BLOCK_Y, 1)
        self.xy = torch.zeros((BLOCK_X * BLOCK_Y, 2), dtype=torch.int)
        for i in range(BLOCK_X):
            self.xy[i * BLOCK_X:(i + 1) * BLOCK_X] = torch.cat(
                (torch.Tensor([x[i]] * BLOCK_Y).int().view(BLOCK_Y, 1), y), 1).int()

    def GetRanges(self, grid, points_xy, depths, radii):
        T1 = time.perf_counter()
        _, indices_min = torch.sort(self.rect_min, dim=0)
        _, indices_max = torch.sort(self.rect_max, dim=0)
        indices_min = indices_min.to(torch.int32)
        indices_max = indices_max.to(torch.int32)
        res = torch.full((grid[0] * grid[1], 50000), -1).to(torch.int32)
        res_size = torch.zeros(grid[0] * grid[1]).to(torch.int32)
        # 找到该tile对应哪些高斯，再对这些高斯的深度进行排序，进而找到每个tile的ranges，以用于渲染
        # pool = ThreadPoolExecutor(max_workers=20)
        # for tile_id in range(grid[0] * grid[1]):
        #     args = (tile_id, grid.data_ptr(), depths.data_ptr(), self.rect_min.data_ptr(),self.rect_max.data_ptr(), self.rect_min.shape[0],
        #             indices_min.data_ptr(), indices_max.data_ptr(), res.data_ptr(), res_size.data_ptr())
        #     pool.submit(GetRanges_tile2, *args)
        # pool.shutdown()
        args = (0, grid.data_ptr(), depths.data_ptr(), self.rect_min.data_ptr(),self.rect_max.data_ptr(), self.rect_min.shape[0],
                indices_min.data_ptr(), indices_max.data_ptr(), res.data_ptr(), res_size.data_ptr())
        GetRanges_tile2(*args)
        for i in range(grid[0] * grid[1]):
            self.ranges.append([res[i][:res_size[i]].int()])
        self.non_empty_tile_id = torch.where(res_size > 0)[0].to(torch.int32)
        # self.non_empty_tile_id = res_size.sort(descending=True).indices[:((res_size > 0).sum())].to(torch.int32)
        # self.non_empty_tile_id = torch.Tensor(self.non_empty_tile_id).to(torch.int32)
        del res
        del res_size
        T2 = time.perf_counter()
        print(f'GetRanges运行时间:{(T2 - T1)*1000}ms')

    def render(self, grid, conic_opacity, W, H, bg, rgb, means2D):
        pool = ThreadPoolExecutor(max_workers=20)
        # grid = torch.tensor(grid).to(torch.int32)
        T1 = time.perf_counter()
        # for n in trange(len(self.non_empty_tile_id)):
        #     tile_id = self.non_empty_tile_id[n]
        #     args=(tile_id, bg.data_ptr(), grid.data_ptr(), means2D.data_ptr(), conic_opacity.data_ptr(), rgb.data_ptr(), 
        #     self.ranges[tile_id][0].data_ptr(), self.ranges[tile_id][0].shape[0],self.image.data_ptr(), H, W)
        #     pool.submit(render_tile, *args)
        # pool.shutdown()

        ranges_address = [self.ranges[i][0].data_ptr() for i in self.non_empty_tile_id]
        ranges_size = [self.ranges[i][0].shape[0] for i in self.non_empty_tile_id]  
        num = self.non_empty_tile_id.shape[0]
        pool.map(render_tile, self.non_empty_tile_id.tolist(), [bg.data_ptr()]*num, [grid.data_ptr()]*num, [means2D.data_ptr()]*num, [conic_opacity.data_ptr()]*num, 
        [rgb.data_ptr()]*num, ranges_address, ranges_size, [self.image.data_ptr()]*num, [H]*num, [W]*num)
        pool.shutdown()

        T2 = time.perf_counter()
        print('render运行时间:%s秒' % (T2 - T1))

    def show(self):
        img = self.image.numpy()[:, :, ::-1]
        cv2.imshow("img", img)

    def Watch(self):
        while self.total > 0:
            self.show()
            if cv2.waitKey(2000) & 0xFF == ord('q'):
                break

class Container(torch.nn.Module):
    def __init__(self, my_values):
        super().__init__()
        for key in my_values:
            setattr(self, key, my_values[key])

def main():
    Input = torch.load("model/ship.pt", map_location=torch.device('cpu'))
    Input = list(Input)
    if type(Input[8]) == tuple:  # 压缩过的，先解压
        centroid = Input[8][0]
        centoird_id = Input[8][1]
        factor = Input[8][2]
        sh_diff = Input[8][3]
        sh_diff = sh_diff / factor
        sh = sh_diff + centroid[centoird_id]
        Input[8] = sh
    with torch.no_grad():
        P = Input[0]
        D = Input[1]
        M = Input[2]
        orig_points = Input[3]
        scales = Input[4]
        scale_modifier = Input[5]
        rotations = Input[6]
        opacities = Input[7]
        shs = Input[8]
        clamped = Input[9]
        cov3D_precomp = Input[10]
        colors_precomp = Input[11]
        viewmatrix = Input[12]
        projmatrix = Input[13]
        cam_pos = Input[14]
        W = Input[15]
        H = Input[16]
        focal_x = Input[17]
        focal_y = Input[18]
        tan_fovx = Input[19]
        tan_fovy = Input[20]
        radii = Input[21]
        means2D = Input[22]
        depths = Input[23]
        cov3Ds = Input[24]
        rgb = Input[25]
        conic_opacity = Input[26]
        tile_grid = Input[27]
        tiles_touched = Input[28]
        prefiltered = Input[29]
        idx = range(0, P)
        mean_time = 0
        max_time = 0
        min_time = 999
        # 预处理
        tile_grid = torch.tensor(tile_grid).int()
        for i in range(0, 360, 2):
            Raster = MyRaster(W, H)
            Raster.rect_min = torch.zeros(P,2).int()
            Raster.rect_max = torch.zeros(P,2).int()
            T1 = time.perf_counter()
            x = i * math.pi / 180
            theta = torch.Tensor([x])
            alpha = torch.Tensor([36])
            gama = torch.Tensor([0])
            z = torch.tensor([[ torch.cos(theta), -torch.sin(theta),  0                ],
                              [ torch.sin(theta),  torch.cos(theta),  0                ],
                              [                0,                 0,  1                ]])
            x = torch.tensor([[                1,                 0,  0                ],
                              [                0,  torch.cos(alpha), -torch.sin(alpha) ],
                              [                0,  torch.sin(alpha),  torch.cos(alpha) ]])
            y = torch.tensor([[  torch.cos(gama),                 0,  torch.sin(gama)  ],
                              [                0,                 1,  0                ],
                              [ -torch.sin(gama),                 0,  torch.cos(gama)  ]])
            R = z @ x @ y
            viewmatrix = torch.tensor([[R[0][0], R[0][1], R[0][2], 0],
                                       [R[1][0], R[1][1], R[1][2], 0],
                                       [R[2][0], R[2][1], R[2][2], 0],
                                       [      0,       0,       4, 1]])
            viewmatrix = torch.tensor([[-4.6591e-02, -1.7530e-01,  9.8341e-01,  0.0000e+00],
                                       [-9.9891e-01,  8.1760e-03, -4.5868e-02,  0.0000e+00],
                                       [-8.1834e-10, -9.8448e-01, -1.7549e-01,  0.0000e+00],
                                       [-3.8155e-09,  1.2222e-09,  4.0311e+00,  1.0000e+00]])
            projection_matrix = torch.tensor([[2.7778, 0.0000, 0.0000, 0.0000],
                                            [0.0000, 2.7778, 0.0000, 0.0000],
                                            [0.0000, 0.0000, 1.0001, 1.0000],
                                            [0.0000, 0.0000, -0.0100, 0.0000]]).to(torch.float)
            projmatrix = viewmatrix @ projection_matrix
            cam_pos = cam_pos.clone()
            T3 = time.perf_counter()
            preprocess(0,
                P,
                D,
                M,
                orig_points.data_ptr(),
                scales.data_ptr(),
                scale_modifier,
                rotations.data_ptr(),
                opacities.data_ptr(),
                shs.data_ptr(),
                clamped.data_ptr(),
                cov3D_precomp.data_ptr(),
                colors_precomp.data_ptr(),
                viewmatrix.data_ptr(),
                projmatrix.data_ptr(),
                cam_pos.data_ptr(),
                W, H,
                tan_fovx, tan_fovy,
                focal_x, focal_y,
                radii.data_ptr(), means2D.data_ptr(),
                depths.data_ptr(), cov3Ds.data_ptr(), rgb.data_ptr(),
                conic_opacity.data_ptr(), tile_grid.data_ptr(),
                tiles_touched.data_ptr(),
                prefiltered,
                Raster.rect_min.data_ptr(),
                Raster.rect_max.data_ptr()
            )
            T4 = time.perf_counter()
            print(f'preprocess运行时间:{(T4 - T3)*1000}ms')
            # 识别每个tile的开始和结束的高斯
            # Raster.GetRanges(tile_grid, means2D, depths, radii)
            # background = torch.tensor([1, 1, 1], dtype=torch.float32)
            # 可视化
            # Raster.total = len(Raster.non_empty_tile_id)
            # p = mp.Process(target=Raster.Watch)
            # p.start()
            # 渲染
            # Raster.render(tile_grid, conic_opacity, W, H, background, rgb, means2D)
            
            #--------------------------------#
            points_offsets = torch.cumsum(tiles_touched,dim=0).int()
            num_rendered = points_offsets[-1]
            gaussian_keys = torch.zeros(num_rendered).long()
            gaussian_values = torch.zeros(num_rendered).int()

            duplicateWithKeys(P,
                            Raster.rect_min.data_ptr(),
                            Raster.rect_max.data_ptr(),
                            depths.data_ptr(),
                            points_offsets.data_ptr(),
                            gaussian_keys.data_ptr(),
                            gaussian_values.data_ptr(),
                            tile_grid.data_ptr())
            indices = gaussian_keys.sort().indices.long() # todo: 改成C
            gaussian_keys = gaussian_keys[indices]
            Raster.gaussian_values = gaussian_values[indices]

            Raster.ranges = torch.zeros((tile_grid[0] * tile_grid[1], 2)).int()
            GetTileRanges(num_rendered, gaussian_keys.data_ptr(), Raster.ranges.data_ptr())
            #--------------------------------#
            render(tile_grid.data_ptr(),
               Raster.gaussian_values.data_ptr(),
               means2D.data_ptr(),
               conic_opacity.data_ptr(),
               rgb.data_ptr(),
               Raster.ranges.data_ptr(),
               Raster.image.data_ptr(),
               H, W)
            #--------------------------------#
            T2 = time.perf_counter()
            print(f'总运行时间:{(T2 - T1)*1000}ms')
            mean_time += (T2-T1)
            max_time = max(T2-T1, max_time)
            min_time = min(T2-T1, min_time)
            image = Raster.image.numpy()
            image = (image * 255).astype(np.uint8)
            image = Image.fromarray(image)
            image.save(f'output/{i}.png', 'png')
            del Raster
            gc.collect()
        mean_time /= 180
        print(f'min_time:{min_time}, max_time:{max_time}, mean_time:{mean_time}')

if __name__ == '__main__':
    # profile = Profile()
    # profile.runcall(main)
    # profile.print_stats(sort='cumtime')
    main()
